use serde::{Deserialize, Serialize};
use std::fs;
use std::path::Path;

use image::ImageReader;
use tauri::{Emitter, Manager};
use std::time::{SystemTime, UNIX_EPOCH};

// PDF rendering imports
use pdfium_render::prelude::*;

pub mod angle_net;
pub mod base_net;
pub mod crnn_net;

pub mod db_net;
pub mod ocr_error;
pub mod ocr_lite;
pub mod ocr_result;
pub mod ocr_utils;
pub mod scale_param;

#[cfg(test)]
mod tests {
    use crate::{ocr_error::OcrError, ocr_lite::OcrLite};
    use std::fs;
    use std::io::{Cursor, Read};

    #[test]
    fn run_test() -> Result<(), OcrError> {
        let mut ocr = OcrLite::new();
        ocr.init_models(
            "./models/ch_PP-OCRv5_mobile_det.onnx",
            "./models/ch_ppocr_mobile_v2.0_cls_infer.onnx",
            "./models/ch_PP-OCRv5_rec_mobile_infer.onnx",
            2,
        )?;

        println!("===test_1===");
        let res = ocr.detect_from_path(
            "./docs/test_images/test_1.png",
            50,
            1024,
            0.5,
            0.3,
            1.6,
            false,
            false,
        )?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });
        println!("===test_2===");
        let res = ocr.detect_from_path(
            "./docs/test_images/test_2.png",
            50,
            1024,
            0.5,
            0.3,
            1.6,
            false,
            false,
        )?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });

        // 通过 image 读取图片
        println!("===test_3===");
        let test_three_img = image::open("./docs/test_images/test_3.png")
            .unwrap()
            .to_rgb8();
        let res = ocr.detect(&test_three_img, 50, 1024, 0.5, 0.3, 1.6, true, false)?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });

        Ok(())
    }

    #[test]
    fn run_test_from_memory() -> Result<(), OcrError> {
        let det_bytes = fs::read("./models/ch_PP-OCRv4_det_infer.onnx")?;
        let cls_bytes = fs::read("./models/ch_ppocr_mobile_v2.0_cls_infer.onnx")?;
        let rec_bytes = fs::read("./models/ch_PP-OCRv4_rec_infer.onnx")?;

        let mut ocr = OcrLite::new();
        ocr.init_models_from_memory(&det_bytes, &cls_bytes, &rec_bytes, 2)?;

        println!("===test_from_memory===");
        let test_img = image::open("./docs/test_images/test_1.png")
            .unwrap()
            .to_rgb8();
        let res = ocr.detect(&test_img, 50, 1024, 0.5, 0.3, 1.6, false, false)?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });

        Ok(())
    }

    #[test]
    fn run_test_from_cursor() -> Result<(), OcrError> {
        let mut det_file = fs::File::open("./models/ch_PP-OCRv4_det_infer.onnx")?;
        let mut cls_file = fs::File::open("./models/ch_ppocr_mobile_v2.0_cls_infer.onnx")?;
        let mut rec_file = fs::File::open("./models/ch_PP-OCRv4_rec_infer.onnx")?;

        let mut det_buffer = Vec::new();
        let mut cls_buffer = Vec::new();
        let mut rec_buffer = Vec::new();

        det_file.read_to_end(&mut det_buffer)?;
        cls_file.read_to_end(&mut cls_buffer)?;
        rec_file.read_to_end(&mut rec_buffer)?;

        let det_cursor = Cursor::new(det_buffer);
        let cls_cursor = Cursor::new(cls_buffer);
        let rec_cursor = Cursor::new(rec_buffer);

        let det_bytes = det_cursor.into_inner();
        let cls_bytes = cls_cursor.into_inner();
        let rec_bytes = rec_cursor.into_inner();

        let mut ocr = OcrLite::new();
        ocr.init_models_from_memory(&det_bytes, &cls_bytes, &rec_bytes, 2)?;

        println!("===test_from_cursor===");
        let test_img = image::open("./docs/test_images/test_2.png")
            .unwrap()
            .to_rgb8();
        let res = ocr.detect(&test_img, 50, 1024, 0.5, 0.3, 1.6, false, false)?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });

        Ok(())
    }

    #[test]
    fn run_test_angle_rollback() -> Result<(), OcrError> {
        let mut ocr = OcrLite::new();
        ocr.init_models(
            "./models/ch_PP-OCRv4_det_infer.onnx",
            "./models/ch_ppocr_mobile_v2.0_cls_infer.onnx",
            "./models/ch_PP-OCRv4_rec_infer.onnx",
            2,
        )?;

        println!("===test_angle_ori===");
        let test_img = image::open("./docs/test_images/test_4.png")
            .unwrap()
            .to_rgb8();
        let res = ocr.detect(&test_img, 50, 1024, 0.5, 0.3, 1.6, true, false)?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });

        println!("===test_angle_rollback===");
        let test_img = image::open("./docs/test_images/test_4.png")
            .unwrap()
            .to_rgb8();
        let res =
            ocr.detect_angle_rollback(&test_img, 50, 1024, 0.5, 0.3, 1.6, true, false, 0.8)?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });

        Ok(())
    }

    #[test]
    fn run_test_from_custom() -> Result<(), OcrError> {
        let mut ocr = OcrLite::new();
        ocr.init_models_custom(
            "./models/ch_PP-OCRv5_mobile_det.onnx",
            "./models/ch_ppocr_mobile_v2.0_cls_infer.onnx",
            "./models/ch_PP-OCRv5_rec_mobile_infer.onnx",
            |builder| Ok(builder.with_inter_threads(2)?.with_intra_threads(2)?),
        )?;

        println!("===test_from_custom===");
        let res = ocr.detect_from_path(
            "./docs/test_images/test_4.png",
            50,
            1024,
            0.5,
            0.3,
            1.6,
            false,
            false,
        )?;
        res.text_blocks.iter().for_each(|item| {
            println!("text: {} score: {}", item.text, item.text_score);
        });

        Ok(())
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct OcrOptions {
  pub model_path: Option<String>,
  pub language: Option<String>,
  pub num_thread: Option<usize>,
  pub padding: Option<u32>,
  pub max_side_len: Option<u32>,
  pub box_score_thresh: Option<f32>,
  pub box_thresh: Option<f32>,
  pub un_clip_ratio: Option<f32>,
  pub do_angle: Option<bool>,
  pub most_angle: Option<bool>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct OcrResult {
  text: String,
  boxes: Option<Vec<Vec<Vec<f32>>>>,
  confidence: Option<Vec<f32>>,
}

#[tauri::command]
async fn ocr_with_paddle(
  image_data: Vec<u8>,
  image_name: String,
  options: OcrOptions,
) -> Result<OcrResult, String> {
  // 创建临时文件保存图片
  let temp_dir = tempfile::tempdir().map_err(|e| format!("创建临时目录失败: {}", e))?;
  let temp_path = temp_dir.path().join(&image_name);
  
  // 保存图片到临时文件
  fs::write(&temp_path, &image_data).map_err(|e| format!("保存图片失败: {}", e))?;
  
  // 加载图片
  let img = ImageReader::open(&temp_path)
    .map_err(|e| format!("加载图片失败: {}", e))?
    .decode()
    .map_err(|e| format!("解码图片失败: {}", e))?;
  
  // 使用默认模型路径或用户指定路径
  let model_path = options.model_path.unwrap_or_else(|| "./models".to_string());
  println!("model_path: {}", model_path);
  // 根据语言选择模型配置
  let lang = options.language.unwrap_or_else(|| "ch".to_string());
  
  // 检查模型文件是否存在
  let det_model = format!("{}/ch_PP-OCRv4_det_infer.onnx", model_path);
  let cls_model = format!("{}/ch_ppocr_mobile_v2.0_cls_infer.onnx", model_path);
  let rec_model = format!("{}/ch_PP-OCRv4_rec_infer.onnx", model_path);
  let keys_file = format!("{}/ppocr_keys.txt", model_path);
  
  if !Path::new(&det_model).exists() {
    return Err(format!("检测模型文件不存在: {}", det_model));
  }
  if !Path::new(&cls_model).exists() {
    return Err(format!("分类模型文件不存在: {}", cls_model));
  }
  if !Path::new(&rec_model).exists() {
    return Err(format!("识别模型文件不存在: {}", rec_model));
  }
  if !Path::new(&keys_file).exists() {
    return Err(format!("字符集文件不存在: {}", keys_file));
  }
  
  // 初始化OCR引擎
  let mut ocr = crate::ocr_lite::OcrLite::new();
  
  // 根据语言加载模型
  let _result = match lang.as_str() {
    "en" => {
      let det_model_en = format!("{}/en_PP-OCRv3_det_infer.onnx", model_path);
      let rec_model_en = format!("{}/en_PP-OCRv3_rec_infer.onnx", model_path);
      
      if Path::new(&det_model_en).exists() && Path::new(&rec_model_en).exists() {
        ocr.init_models(&det_model_en, &cls_model, &rec_model_en, options.num_thread.unwrap_or(2))
          .map_err(|e| format!("初始化英文OCR模型失败: {}", e))?
      } else {
        ocr.init_models(&det_model, &cls_model, &rec_model, options.num_thread.unwrap_or(2))
          .map_err(|e| format!("初始化默认OCR模型失败: {}", e))?
      }
    }
    _ => {
      ocr.init_models(&det_model, &cls_model, &rec_model, options.num_thread.unwrap_or(2))
        .map_err(|e| format!("初始化中文OCR模型失败: {}", e))?
    }
  };
  
  // 转换图片为RGB格式
  let rgb_img = img.to_rgb8();
  
  // 执行OCR检测
  let ocr_result = ocr.detect(
    &rgb_img,
    options.padding.unwrap_or(50),
    options.max_side_len.unwrap_or(1024),
    options.box_score_thresh.unwrap_or(0.5),
    options.box_thresh.unwrap_or(0.3),
    options.un_clip_ratio.unwrap_or(1.6),
    options.do_angle.unwrap_or(true),
    options.most_angle.unwrap_or(false),
  ).map_err(|e| format!("OCR检测失败: {}", e))?;
  
  // 提取文本和置信度
  let mut full_text = String::new();
  let mut boxes = Vec::new();
  let mut confidences = Vec::new();
  
  for text_block in &ocr_result.text_blocks {
    if !full_text.is_empty() {
      full_text.push('\n');
    }
    full_text.push_str(&text_block.text);
    
    // 转换box坐标为Vec<Vec<f32>>格式
    let mut box_coords = Vec::new();
    for point in &text_block.box_points {
      let mut coord = Vec::new();
      coord.push(point.x as f32);
      coord.push(point.y as f32);
      box_coords.push(coord);
    }
    boxes.push(box_coords);
    
    // 添加置信度
    confidences.push(text_block.text_score as f32);
  }
  
  // 保存图片和识别结果
  let timestamp = SystemTime::now()
    .duration_since(UNIX_EPOCH)
    .map_err(|e| format!("获取时间戳失败: {}", e))?
    .as_secs();
    
  // 移除了保存到文件的代码
  
  // 移除了保存到数据库的代码
  
  Ok(OcrResult {
    text: full_text,
    boxes: Some(boxes),
    confidence: Some(confidences),
  })
}

#[tauri::command]
fn show_main_window(app: tauri::AppHandle) {
    if let Some(window) = app.get_webview_window("main") {
        let _ = window.show();
        let _ = window.set_focus();
    }
}

#[tauri::command]
fn open_settings(app: tauri::AppHandle) {
    if let Some(window) = app.get_webview_window("main") {
        let _ = window.show();
        let _ = window.set_focus();
        let _ = window.eval("window.location.href = '/settings'");
    }
}



// 添加PDF处理函数
#[tauri::command]
async fn ocr_pdf_with_paddle(
  pdf_data: Vec<u8>,
  pdf_name: String,
  options: OcrOptions,
) -> Result<OcrResult, String> {
  // 创建临时目录保存PDF转换的图片
  let temp_dir = tempfile::tempdir().map_err(|e| format!("创建临时目录失败: {}", e))?;
  let temp_path = temp_dir.path().join(&pdf_name);
  
  // 保存PDF到临时文件
  fs::write(&temp_path, &pdf_data).map_err(|e| format!("保存PDF失败: {}", e))?;
  
  let mut full_text = String::new();
  let mut all_boxes = Vec::new();
  let mut all_confidences = Vec::new();
  
  // 初始化OCR引擎
  let mut ocr = crate::ocr_lite::OcrLite::new();
  
  // 使用默认模型路径或用户指定路径
  let model_path = options.model_path.unwrap_or_else(|| "./models".to_string());
  
  // 根据语言选择模型配置
  let lang = options.language.unwrap_or_else(|| "ch".to_string());
  
  // 检查模型文件是否存在
  let det_model = format!("{}/ch_PP-OCRv4_det_infer.onnx", model_path);
  let cls_model = format!("{}/ch_ppocr_mobile_v2.0_cls_infer.onnx", model_path);
  let rec_model = format!("{}/ch_PP-OCRv4_rec_infer.onnx", model_path);
  let keys_file = format!("{}/ppocr_keys.txt", model_path);
  
  if !Path::new(&det_model).exists() {
    return Err(format!("检测模型文件不存在: {}", det_model));
  }
  if !Path::new(&cls_model).exists() {
    return Err(format!("分类模型文件不存在: {}", cls_model));
  }
  if !Path::new(&rec_model).exists() {
    return Err(format!("识别模型文件不存在: {}", rec_model));
  }
  if !Path::new(&keys_file).exists() {
    return Err(format!("字符集文件不存在: {}", keys_file));
  }
  
  // 根据语言加载模型
  let _result = match lang.as_str() {
    "en" => {
      let det_model_en = format!("{}/en_PP-OCRv3_det_infer.onnx", model_path);
      let rec_model_en = format!("{}/en_PP-OCRv3_rec_infer.onnx", model_path);
      
      if Path::new(&det_model_en).exists() && Path::new(&rec_model_en).exists() {
        ocr.init_models(&det_model_en, &cls_model, &rec_model_en, options.num_thread.unwrap_or(2))
          .map_err(|e| format!("初始化英文OCR模型失败: {}", e))?
      } else {
        ocr.init_models(&det_model, &cls_model, &rec_model, options.num_thread.unwrap_or(2))
          .map_err(|e| format!("初始化默认OCR模型失败: {}", e))?
      }
    }
    _ => {
      ocr.init_models(&det_model, &cls_model, &rec_model, options.num_thread.unwrap_or(2))
        .map_err(|e| format!("初始化中文OCR模型失败: {}", e))?
    }
  };
  
  // Load PDF document using pdfium-render
  //let pdfium = Pdfium::default();
  let pdfium = Pdfium::new(
        Pdfium::bind_to_library(Pdfium::pdfium_platform_library_name_at_path("./"))
            .or_else(|_| Pdfium::bind_to_system_library())
            .unwrap() // Or use the ? unwrapping operator to pass any error up to the caller
    );
  let document = pdfium.load_pdf_from_file(&temp_path, None)
    .map_err(|e| format!("加载PDF失败: {}", e))?;

  // 遍历PDF的每一页
  for page_index in 0..document.pages().len() {
    let page = document.pages().get(page_index).map_err(|e| format!("获取PDF页面失败: {}", e))?;

    // 渲染页面为图片
    let render_config = PdfRenderConfig::new()
      .set_target_width(2000)  // 设置目标宽度
      .set_maximum_height(2000)  // 设置最大高度
      .rotate_if_landscape(PdfPageRenderRotation::Degrees90, true);  // 如果是横向页面则旋转
    
    let image = page.render_with_config(&render_config)
      .map_err(|e| format!("渲染PDF页面失败: {}", e))?
      .as_image();  // 转换为image::DynamicImage
    
    // 创建临时文件保存渲染后的图片
    let image_path = temp_dir.path().join(format!("page_{}.jpg", page_index));
    image.save_with_format(&image_path, image::ImageFormat::Jpeg)
      .map_err(|e| format!("保存图片失败: {}", e))?;
    
    // 打印 临时保存图片的路径
    println!("临时保存图片的路径: {:?}", image_path);

    // 对图片进行OCR识别
    let img = ImageReader::open(&image_path)
      .map_err(|e| format!("打开图片失败: {}", e))?
      .decode()
      .map_err(|e| format!("解码图片失败: {}", e))?
      .into_rgb8(); // 转换为RgbImage

    let ocr_result = ocr.detect(
      &img, // 传递引用
      options.padding.unwrap_or(50),
      options.max_side_len.unwrap_or(1024),
      options.box_score_thresh.unwrap_or(0.5),
      options.box_thresh.unwrap_or(0.3),
      options.un_clip_ratio.unwrap_or(1.6),
      options.do_angle.unwrap_or(true),
      options.most_angle.unwrap_or(true),
    ).map_err(|e| format!("OCR识别失败: {}", e))?;

    // 提取文本和置信度
    for text_block in &ocr_result.text_blocks {
      if !full_text.is_empty() {
        full_text.push('\n');
      }
      full_text.push_str(&text_block.text);
      
      // 转换box坐标为Vec<Vec<f32>>格式
      let mut box_coords = Vec::new();
      for point in &text_block.box_points {
        let mut coord = Vec::new();
        coord.push(point.x as f32);
        coord.push(point.y as f32);
        box_coords.push(coord);
      }
      all_boxes.push(box_coords);
      
      // 添加置信度
      all_confidences.push(text_block.text_score as f32);
    }
  }
  
  // 保存PDF和识别结果
  let timestamp = SystemTime::now()
    .duration_since(UNIX_EPOCH)
    .map_err(|e| format!("获取时间戳失败: {}", e))?
    .as_secs();
    
  // 移除了保存到文件的代码
  
  // 移除了保存到数据库的代码
  
  Ok(OcrResult {
    text: full_text,
    boxes: Some(all_boxes),
    confidence: Some(all_confidences),
  })
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    tauri::Builder::default()
        .plugin(tauri_plugin_dialog::init())
        .plugin(tauri_plugin_fs::init())
        .plugin(tauri_plugin_shell::init())
        .plugin(tauri_plugin_updater::Builder::new().build())
        .invoke_handler(tauri::generate_handler![
            ocr_with_paddle, 
            ocr_pdf_with_paddle, // 添加PDF处理函数
            show_main_window, 
            open_settings,

        ])
        .setup(|app| {
            if cfg!(debug_assertions) {
                app.handle().plugin(
                    tauri_plugin_log::Builder::default()
                        .level(log::LevelFilter::Info)
                        .build(),
                )?;
            }

            // 移除了数据库初始化代码

            // 设置系统托盘 - 简化版本，使用默认图标
            use tauri::tray::{TrayIconBuilder, MouseButton, TrayIconEvent};
            use tauri::menu::{Menu, MenuItem, PredefinedMenuItem};
            use tauri_plugin_updater::UpdaterExt;

            // 设置窗口关闭事件监听器
            if let Some(window) = app.get_webview_window("main") {
                let app_handle = app.handle().clone();
                let window_clone = window.clone();
                window.on_window_event(move |event| {
                    if let tauri::WindowEvent::CloseRequested { api, .. } = event {
                        // 阻止窗口关闭，改为隐藏窗口
                        api.prevent_close();
                        let _ = window_clone.hide();
                        
                        // 显示系统托盘通知
                        if let Some(tray) = app_handle.tray_by_id("main") {
                            let _ = tray.set_tooltip(Some("MoreOCR 已最小化到系统托盘"));
                        }
                    }
                });
            }

            let menu = Menu::with_items(app, &[
                &MenuItem::with_id(app, "show", "显示主窗口", true, None::<&str>)?,
                &MenuItem::with_id(app, "screenshot", "截图识别", true, None::<&str>)?,
                &MenuItem::with_id(app, "settings", "设置", true, None::<&str>)?,
                &PredefinedMenuItem::separator(app)?,
                &PredefinedMenuItem::quit(app, Some("退出"))?,
            ])?;

            // 创建托盘图标，使用默认的icon
            let _tray = TrayIconBuilder::with_id("main")
                .icon(app.default_window_icon().unwrap().clone())
                .tooltip("MoreOCR")
                .menu(&menu)
                .on_menu_event(move |app, event| match event.id.as_ref() {
                    "show" => {
                        if let Some(window) = app.get_webview_window("main") {
                            let _ = window.show();
                            let _ = window.set_focus();
                        }
                    }
                    "screenshot" => {
                        if let Some(window) = app.get_webview_window("main") {
                            let _ = window.emit("tray-screenshot", ());
                            let _ = window.show();
                            let _ = window.set_focus();
                        }
                    }
                    "settings" => {
                        if let Some(window) = app.get_webview_window("main") {
                            let _ = window.show();
                            let _ = window.set_focus();
                            let _ = window.eval("window.location.href = '/settings'");
                        }
                    }
                    _ => {}
                })
                .on_tray_icon_event(|tray, event| {
                    if let TrayIconEvent::Click {
                        button: MouseButton::Left,
                        button_state: _,
                        .. 
                    } = event {
                        println!("系统托盘左键点击");
                        let app = tray.app_handle();
                        if let Some(window) = app.get_webview_window("main") {
                            let _ = window.show();
                            let _ = window.set_focus();
                        }
                    }
                })
                .build(app)?;

            // 启动时检查更新（延迟5秒，避免影响启动速度）
            let handle = app.handle().clone();
            tauri::async_runtime::spawn(async move {
                tokio::time::sleep(std::time::Duration::from_secs(5)).await;
                if let Ok(updater) = handle.updater() {
                    let _ = updater.check().await;
                }
            });

            Ok(())
        })
    .run(tauri::generate_context!())
    .expect("error while running tauri application");
}
